//
// Created by 1234 on 2025/3/6.
//

#include "VkGPUHelper.h"

#include <iostream>
#include <ostream>
#ifdef OS_OPEN_HARMONY
#include <effect_engine/gpu/utils/vk_enum_string_helper.h>
#else
#include <vulkan/vk_enum_string_helper.h>
#endif

#include "effect_engine/log/Log.h"
#include "effect_engine/utils/IOUtils.h"
#include "effect_engine/utils/TimeUtils.h"

VkResult VkGPUHelper::CreateUniformBufferAndUploadData(const VkDevice device,
                                                       const std::vector<uint32_t> &queueFamilyIndices,
                                                       const VkPhysicalDeviceMemoryProperties *memoryProperties,
                                                       const VkDeviceSize bufferSize,
                                                       VkBuffer *buffer,
                                                       VkDeviceMemory *bufferMemory,
                                                       const void *uploadData) {
    VkResult ret = CreateStorageBufferAndBindMem(device,
                                                 bufferSize,
                                                 queueFamilyIndices,
                                                 memoryProperties,
                                                 buffer,
                                                 bufferMemory);
    if (ret != VK_SUCCESS) {
        Logger() << "Failed to create input storage buffer!" << std::endl;
        return ret;
    }

    void *data = nullptr;
    ret = vkMapMemory(device, *bufferMemory, 0, bufferSize, 0, &data);
    if (ret != VK_SUCCESS) {
        Logger() << "Failed to map input storage buffer memory, err=" << string_VkResult(ret) << std::endl;
        return ret;
    }
    memcpy(data, uploadData, bufferSize);
    vkUnmapMemory(device, *bufferMemory);
    return ret;
}

VkResult VkGPUHelper::CreateStorageBufferAndUploadData(const VkDevice device,
                                                       const std::vector<uint32_t> &queueFamilyIndices,
                                                       const VkPhysicalDeviceMemoryProperties *memoryProperties,
                                                       const VkDeviceSize bufferSize,
                                                       VkBuffer *buffer,
                                                       VkDeviceMemory *bufferMemory,
                                                       const void *uploadData) {
    VkResult ret = CreateStorageBufferAndBindMem(device,
                                                 bufferSize,
                                                 queueFamilyIndices,
                                                 memoryProperties,
                                                 buffer,
                                                 bufferMemory);
    if (ret != VK_SUCCESS) {
        Logger() << "Failed to create input storage buffer!" << std::endl;
        return ret;
    }

    void *data = nullptr;
    ret = vkMapMemory(device, *bufferMemory, 0, bufferSize, 0, &data);
    if (ret != VK_SUCCESS) {
        Logger() << "Failed to map input storage buffer memory, err=" << string_VkResult(ret) << std::endl;
        return ret;
    }
    memcpy(data, uploadData, bufferSize);
    vkUnmapMemory(device, *bufferMemory);
    return ret;
}

VkSubmitInfo VkGPUHelper::BuildSubmitInfo(const VkFlags *submitWaitDstStageMask,
                                          const std::vector<VkCommandBuffer> &submitCommandBuffers,
                                          const std::vector<VkSemaphore> &submitSignalSemaphores,
                                          const std::vector<VkSemaphore> &submitWaitSemaphores) {
    VkSubmitInfo submitInfo = {};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submitInfo.pNext = nullptr;
    submitInfo.commandBufferCount = submitCommandBuffers.size();
    submitInfo.pCommandBuffers = submitCommandBuffers.data();
    submitInfo.signalSemaphoreCount = submitSignalSemaphores.size();
    submitInfo.pSignalSemaphores = submitSignalSemaphores.data();
    submitInfo.waitSemaphoreCount = submitWaitSemaphores.size();
    submitInfo.pWaitSemaphores = submitWaitSemaphores.data();
    submitInfo.pWaitDstStageMask = submitWaitDstStageMask;
    return submitInfo;
}

VkBufferMemoryBarrier VkGPUHelper::BuildBufferMemoryBarrier(const VkAccessFlagBits srcAccessMask,
                                                            const VkAccessFlagBits dstAccessMask,
                                                            const VkBuffer buffer,
                                                            const VkDeviceSize size) {
    VkBufferMemoryBarrier bufferMemoryBarrier = {};
    bufferMemoryBarrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
    bufferMemoryBarrier.pNext = nullptr;
    bufferMemoryBarrier.srcAccessMask = srcAccessMask;
    bufferMemoryBarrier.dstAccessMask = dstAccessMask;
    bufferMemoryBarrier.buffer = buffer;
    bufferMemoryBarrier.offset = 0;
    bufferMemoryBarrier.size = size;
    bufferMemoryBarrier.dstQueueFamilyIndex = 0;
    bufferMemoryBarrier.srcQueueFamilyIndex = 0;
    return bufferMemoryBarrier;
}

void VkGPUHelper::GPUCmdPipelineBufferMemBarrier(const VkCommandBuffer commandBuffer,
                                                 const VkPipelineStageFlags srcStageMask,
                                                 const VkPipelineStageFlags dstStageMask,
                                                 const VkDependencyFlags dependencyFlags,
                                                 const std::vector<VkBufferMemoryBarrier> &bufferMemoryBarriers) {
    const std::vector<VkMemoryBarrier> memoryBarriers;
    const std::vector<VkImageMemoryBarrier> imageMemoryBarrier;
    GPUCmdPipelineBarrier(commandBuffer,
                          srcStageMask,
                          dstStageMask,
                          dependencyFlags,
                          memoryBarriers,
                          bufferMemoryBarriers,
                          imageMemoryBarrier);
}

void VkGPUHelper::GPUCmdPipelineBarrier(const VkCommandBuffer commandBuffer,
                                        const VkPipelineStageFlags srcStageMask,
                                        const VkPipelineStageFlags dstStageMask,
                                        const VkDependencyFlags dependencyFlags,
                                        const std::vector<VkMemoryBarrier> &memoryBarriers,
                                        const std::vector<VkBufferMemoryBarrier> &bufferMemoryBarriers,
                                        const std::vector<VkImageMemoryBarrier> &imageMemoryBarriers
) {
    vkCmdPipelineBarrier(commandBuffer,
                         srcStageMask,
                         dstStageMask,
                         dependencyFlags,
                         memoryBarriers.size(),
                         memoryBarriers.data(),
                         bufferMemoryBarriers.size(),
                         bufferMemoryBarriers.data(),
                         imageMemoryBarriers.size(),
                         imageMemoryBarriers.data());
}

VkResult VkGPUHelper::GPUQueueSubmit(const VkQueue queue,
                                     const std::vector<VkSubmitInfo> &submitInfos,
                                     const VkFence fence) {
    const VkResult ret = vkQueueSubmit(queue, submitInfos.size(), submitInfos.data(), fence);
    if (ret != VK_SUCCESS) {
        Logger() << "vkQueueSubmit failed, err=" << string_VkResult(ret) << std::endl;
    }
    return ret;
}

void VkGPUHelper::GPUCmdDispatch(const VkCommandBuffer commandBuffer,
                                 const uint32_t groupCountX,
                                 const uint32_t groupCountY,
                                 const uint32_t groupCountZ) {
    vkCmdDispatch(commandBuffer, groupCountX, groupCountY, groupCountZ);
}

void VkGPUHelper::GPUCmdPushConstant(const VkCommandBuffer commandBuffer,
                                     const VkPipelineLayout pipelineLayout,
                                     VkShaderStageFlagBits shaderStageFlag,
                                     const uint32_t offset,
                                     const size_t pushConstantSize,
                                     const void *pushConstantData) {
    vkCmdPushConstants(commandBuffer,
                       pipelineLayout,
                       VK_SHADER_STAGE_COMPUTE_BIT,
                       offset,
                       pushConstantSize,
                       pushConstantData);
}

void VkGPUHelper::GPUEndCommandBuffer(const VkCommandBuffer commandBuffer) {
    vkEndCommandBuffer(commandBuffer);
}

void VkGPUHelper::GPUBeginCommandBuffer(const VkCommandBuffer commandBuffer) {
    VkCommandBufferBeginInfo beginInfo = {};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    beginInfo.flags = 0;
    beginInfo.pNext = nullptr;
    beginInfo.pInheritanceInfo = nullptr;
    vkBeginCommandBuffer(commandBuffer, &beginInfo);
}

VkResult VkGPUHelper::AllocateCommandBuffers(const VkDevice device,
                                             const VkCommandPool commandPool,
                                             const uint32_t commandBufferCount,
                                             VkCommandBuffer *commandBuffers) {
    VkCommandBufferAllocateInfo allocInfo = {};
    allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocInfo.pNext = nullptr;
    allocInfo.commandPool = commandPool;
    allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocInfo.commandBufferCount = commandBufferCount;

    const VkResult result = vkAllocateCommandBuffers(device, &allocInfo, commandBuffers);
    if (result != VK_SUCCESS) {
        Logger() << "Failed to allocate command buffer, err = " << string_VkResult(result) << std::endl;
    }
    return result;
}

VkResult VkGPUHelper::CreateFence(const VkDevice device, VkFence *fence) {
    VkFenceCreateInfo info = {};
    info.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
    info.flags = VK_FENCE_CREATE_SIGNALED_BIT;
    info.pNext = nullptr;
    const VkResult result = vkCreateFence(device, &info, nullptr, fence);
    if (result != VK_SUCCESS) {
        Logger() << "vkCreateFence failed, err=" << string_VkResult(result) << std::endl;
    }
    return result;
}

VkWriteDescriptorSet VkGPUHelper::BuildWriteDescriptorSet(const VkDescriptorSet descriptorSet,
                                                          const uint32_t dtsBinding,
                                                          const VkDescriptorType type,
                                                          const VkDescriptorImageInfo *imageInfo,
                                                          const VkDescriptorBufferInfo *bufferInfo,
                                                          const VkBufferView *texelBufferView) {
    VkWriteDescriptorSet writeDescriptorSet = {};
    writeDescriptorSet.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writeDescriptorSet.dstSet = descriptorSet;
    writeDescriptorSet.dstBinding = dtsBinding;
    writeDescriptorSet.dstArrayElement = 0;
    writeDescriptorSet.descriptorType = type;
    writeDescriptorSet.descriptorCount = 1;
    writeDescriptorSet.pBufferInfo = bufferInfo;
    writeDescriptorSet.pImageInfo = imageInfo;
    writeDescriptorSet.pTexelBufferView = texelBufferView;
    return writeDescriptorSet;
}

VkWriteDescriptorSet VkGPUHelper::BuildWriteStorageBufferDescriptorSet(const VkDescriptorSet descriptorSet,
                                                                       const uint32_t dstBinding,
                                                                       const VkDescriptorBufferInfo *
                                                                       descriptorBufferInfo) {
    return BuildWriteDescriptorSet(descriptorSet,
                                   dstBinding,
                                   VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
                                   nullptr,
                                   descriptorBufferInfo,
                                   nullptr);
}

VkResult VkGPUHelper::AllocateDescriptorSets(const VkDevice device,
                                             const VkDescriptorPool descriptorPool,
                                             const std::vector<VkDescriptorSetLayout> &descriptorSetLayouts,
                                             VkDescriptorSet *descriptorSets) {
    VkDescriptorSetAllocateInfo descriptorSetAllocateInfo{};
    descriptorSetAllocateInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    descriptorSetAllocateInfo.pNext = nullptr;
    descriptorSetAllocateInfo.descriptorPool = descriptorPool;
    descriptorSetAllocateInfo.descriptorSetCount = descriptorSetLayouts.size();
    descriptorSetAllocateInfo.pSetLayouts = descriptorSetLayouts.data();

    const VkResult result = vkAllocateDescriptorSets(device, &descriptorSetAllocateInfo, descriptorSets);
    if (result != VK_SUCCESS) {
        Logger() << "failed to allocate descriptor sets, err=" << string_VkResult(result) << std::endl;
    }
    return result;
}

VkResult VkGPUHelper::CreateShaderModule(const VkDevice device,
                                         const size_t shaderCodeSize,
                                         const uint32_t *shaderSpvCode,
                                         VkShaderModule *shaderModule) {
    VkShaderModuleCreateInfo createInfo = {};
    createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
    createInfo.flags = 0;
    createInfo.pNext = nullptr;
    createInfo.codeSize = shaderCodeSize;
    createInfo.pCode = shaderSpvCode;

    const VkResult result = vkCreateShaderModule(device, &createInfo, nullptr, shaderModule);
    if (result != VK_SUCCESS) {
        Logger() << "failed to create shader module, err=" << string_VkResult(result) << std::endl;
    }
    return result;
}

VkResult VkGPUHelper::CreateShaderModuleFromPath(const VkDevice device,
                                                 const std::string &shaderFilePath,
                                                 VkShaderModule *shaderModule) {
    std::vector<char> shaderSpvCode = IOUtils::ReadFile(shaderFilePath);
    if (shaderSpvCode.empty()) {
        return VK_ERROR_INVALID_SHADER_NV;
    }
    const VkResult result = CreateShaderModule(device,
                                               shaderSpvCode.size(),
                                               reinterpret_cast<uint32_t *>(shaderSpvCode.data()),
                                               shaderModule);
    if (result != VK_SUCCESS) {
        Logger() << "Failed to create shader module, err=" << string_VkResult(result) << std::endl;
    }
    return result;
}

VkResult VkGPUHelper::CreateComputePipeline(const VkDevice device,
                                            const VkPipelineCache pipelineCache,
                                            const VkPipelineLayout pipelineLayout,
                                            const VkShaderModule shaderModule,
                                            VkPipeline *pipeline) {
    VkPipelineShaderStageCreateInfo shaderStageInfo = {};
    shaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    shaderStageInfo.flags = 0;
    shaderStageInfo.pNext = nullptr;
    shaderStageInfo.pName = "main";
    shaderStageInfo.stage = VK_SHADER_STAGE_COMPUTE_BIT;
    shaderStageInfo.module = shaderModule;
    shaderStageInfo.pSpecializationInfo = nullptr;

    std::vector<VkComputePipelineCreateInfo> pipelineCreateInfos;
    VkComputePipelineCreateInfo computePipelineCreateInfo;
    computePipelineCreateInfo.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
    computePipelineCreateInfo.flags = 0;
    computePipelineCreateInfo.pNext = nullptr;
    computePipelineCreateInfo.basePipelineIndex = 0;
    computePipelineCreateInfo.basePipelineHandle = VK_NULL_HANDLE;
    computePipelineCreateInfo.layout = pipelineLayout;
    computePipelineCreateInfo.stage = shaderStageInfo;
    pipelineCreateInfos.push_back(computePipelineCreateInfo);
    const VkResult result = vkCreateComputePipelines(device,
                                                     pipelineCache,
                                                     pipelineCreateInfos.size(),
                                                     pipelineCreateInfos.data(),
                                                     nullptr,
                                                     pipeline);
    if (result != VK_SUCCESS) {
        Logger() << "failed to create compute pipeline, err=" << string_VkResult(result) << std::endl;
    }
    return result;
}

VkResult VkGPUHelper::CreatePipelineLayout(const VkDevice device,
                                           const std::vector<VkDescriptorSetLayout> &descriptorSetLayouts,
                                           const std::vector<VkPushConstantRange> &pushConstantRanges,
                                           VkPipelineLayout *pipelineLayout) {
    VkPipelineLayoutCreateInfo pipelineLayoutInfo = {};
    pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutInfo.flags = 0;
    pipelineLayoutInfo.pNext = nullptr;
    pipelineLayoutInfo.setLayoutCount = descriptorSetLayouts.size();
    pipelineLayoutInfo.pSetLayouts = descriptorSetLayouts.data();
    pipelineLayoutInfo.pushConstantRangeCount = pushConstantRanges.size();
    pipelineLayoutInfo.pPushConstantRanges = pushConstantRanges.data();

    const VkResult ret = vkCreatePipelineLayout(device, &pipelineLayoutInfo, nullptr, pipelineLayout);
    if (ret != VK_SUCCESS) {
        Logger() << "failed to create pipeline layout, " << string_VkResult(ret) << std::endl;
    }
    return ret;
}

VkResult VkGPUHelper::CreateDescriptorSetLayout(const VkDevice device,
                                                const std::vector<VkDescriptorSetLayoutBinding> &
                                                descriptorSetLayoutBindings,
                                                VkDescriptorSetLayout *descriptorSetLayout) {
    VkDescriptorSetLayoutCreateInfo descriptorSetLayoutCreateInfo;
    descriptorSetLayoutCreateInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    descriptorSetLayoutCreateInfo.flags = 0;
    descriptorSetLayoutCreateInfo.pNext = nullptr;
    descriptorSetLayoutCreateInfo.bindingCount = descriptorSetLayoutBindings.size();
    descriptorSetLayoutCreateInfo.pBindings = descriptorSetLayoutBindings.data();

    const VkResult ret = vkCreateDescriptorSetLayout(device,
                                                     &descriptorSetLayoutCreateInfo,
                                                     nullptr,
                                                     descriptorSetLayout);
    if (ret != VK_SUCCESS) {
        std::cerr << "failed to create descriptor set layout, err=" << string_VkResult(ret) << std::endl;
        return ret;
    }
    return ret;
}

VkResult VkGPUHelper::CreateBuffer(const VkDevice device,
                                   const VkDeviceSize size,
                                   const VkBufferUsageFlagBits usage,
                                   const VkSharingMode sharingMode,
                                   const std::vector<uint32_t> &queueFamilyIndices,
                                   VkBuffer *buffer) {
    VkBufferCreateInfo bufferInfo = {};
    bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
    bufferInfo.flags = 0;
    bufferInfo.pNext = nullptr;
    bufferInfo.size = size;
    bufferInfo.usage = usage;
    bufferInfo.sharingMode = sharingMode;
    bufferInfo.queueFamilyIndexCount = queueFamilyIndices.size();
    bufferInfo.pQueueFamilyIndices = queueFamilyIndices.data();
    const VkResult ret = vkCreateBuffer(device, &bufferInfo, nullptr, buffer);
    if (ret != VK_SUCCESS) {
        Logger() << "vkCreateBuffer failed, err=" << string_VkResult(ret) << std::endl;
    }
    return ret;
}

uint32_t VkGPUHelper::GetRequiredMemTypeIndex(const VkPhysicalDeviceMemoryProperties *memProps,
                                              const VkMemoryRequirements &memRequirements,
                                              const uint32_t memoryPropertyFlagsBits) {
    for (uint32_t i = 0; i < memProps->memoryTypeCount; i++) {
        if (memRequirements.memoryTypeBits & (1 << i)) {
            if ((memoryPropertyFlagsBits & memProps->memoryTypes[i].propertyFlags) == memoryPropertyFlagsBits) {
                return i;
            }
        }
    }
    return 0;
}

VkResult VkGPUHelper::CreateStorageBufferAndBindMem(const VkDevice device,
                                                    const VkDeviceSize size,
                                                    const std::vector<uint32_t> &queueFamilyIndices,
                                                    const VkPhysicalDeviceMemoryProperties *memProps,
                                                    VkBuffer *storageBuffer,
                                                    VkDeviceMemory *storageBufferMemory) {
    VkResult ret = CreateBuffer(device,
                                size,
                                VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
                                VK_SHARING_MODE_EXCLUSIVE,
                                queueFamilyIndices,
                                storageBuffer);
    if (ret != VK_SUCCESS) {
        Logger() << "vkCreateBuffer failed, err=" << string_VkResult(ret) << std::endl;
    }

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(device, *storageBuffer, &memRequirements);

    VkMemoryAllocateInfo allocInfo = {};
    allocInfo.pNext = nullptr;
    allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.allocationSize = memRequirements.size;
    allocInfo.memoryTypeIndex = GetRequiredMemTypeIndex(memProps,
                                                        memRequirements,
                                                        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
                                                        VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
    ret = vkAllocateMemory(device, &allocInfo, nullptr, storageBufferMemory);
    if (ret != VK_SUCCESS) {
        Logger() << "vkAllocateMemory failed, err=" << string_VkResult(ret) << std::endl;
        return ret;
    }
    ret = vkBindBufferMemory(device, *storageBuffer, *storageBufferMemory, 0);
    if (ret != VK_SUCCESS) {
        Logger() << "vkBindBufferMemory failed, err=" << string_VkResult(ret) << std::endl;
        return ret;
    }
    return ret;
}


VkResult VkGPUHelper::CreateUniformBufferAndBindMem(const VkDevice device,
                                                    const VkDeviceSize size,
                                                    const std::vector<uint32_t> &queueFamilyIndices,
                                                    const VkPhysicalDeviceMemoryProperties *memProps,
                                                    VkBuffer *storageBuffer,
                                                    VkDeviceMemory *storageBufferMemory) {
    VkResult ret = CreateBuffer(device,
                                size,
                                VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
                                VK_SHARING_MODE_EXCLUSIVE,
                                queueFamilyIndices,
                                storageBuffer);
    if (ret != VK_SUCCESS) {
        Logger() << "vkCreateBuffer failed, err=" << string_VkResult(ret) << std::endl;
    }

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(device, *storageBuffer, &memRequirements);

    VkMemoryAllocateInfo allocInfo = {};
    allocInfo.pNext = nullptr;
    allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.allocationSize = memRequirements.size;
    allocInfo.memoryTypeIndex = GetRequiredMemTypeIndex(memProps,
                                                        memRequirements,
                                                        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
                                                        VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
    ret = vkAllocateMemory(device, &allocInfo, nullptr, storageBufferMemory);
    if (ret != VK_SUCCESS) {
        Logger() << "vkAllocateMemory failed, err=" << string_VkResult(ret) << std::endl;
        return ret;
    }
    ret = vkBindBufferMemory(device, *storageBuffer, *storageBufferMemory, 0);
    if (ret != VK_SUCCESS) {
        Logger() << "vkBindBufferMemory failed, err=" << string_VkResult(ret) << std::endl;
        return ret;
    }
    return ret;
}
